import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader, random_split

# 数据预处理和加载
def load_data(batch_size=64):
    # 定义数据预处理流程：将图片转换为Tensor，并进行归一化处理
    transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])

    # 加载MNIST数据集
    train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
    test_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)

    # 使用DataLoader来加载数据
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

    return train_loader, test_loader

# BP神经网络结构
class BPNet(nn.Module):
    def __init__(self, input_size=784, hidden_size=128, output_size=10):
        super(BPNet, self).__init__()
        # 定义网络层：输入层 -> 隐藏层 -> 输出层
        self.fc1 = nn.Linear(input_size, hidden_size)  # 输入层到隐藏层
        self.fc2 = nn.Linear(hidden_size, output_size)  # 隐藏层到输出层
        self.relu = nn.ReLU()  # 激活函数ReLU

    def forward(self, x):
        # 前向传播
        x = x.view(-1, 784)  # 将每个28x28的图像展平成784维向量
        x = self.relu(self.fc1(x))  # 输入通过第一个线性层和ReLU激活函数
        x = self.fc2(x)  # 输入通过第二个线性层
        return x

# 权重初始化函数
def init_weights(m):
    if isinstance(m, nn.Linear):
        nn.init.normal_(m.weight, mean=0, std=0.01)  # 权重初始化为正态分布
        nn.init.constant_(m.bias, 0)  # 偏置初始化为0

# 创建并初始化网络
model = BPNet()
model.apply(init_weights)

# 损失函数和优化器
criterion = nn.CrossEntropyLoss()  # 交叉熵损失
optimizer = optim.SGD(model.parameters(), lr=0.01)  # 随机梯度下降优化器
